package luke.audio.analyzer;

import java.util.*;

/**
 * BPM计算器类
 * 提供多种BPM计算方法，包括自相关算法和基于频带能量的分析
 */
public class BPMCalculator {


    /**
     * 使用自相关方法计算BPM（改进版）
     * 1. 计算自相关并找到峰值
     * 2. 计算5-250 BPM对应的时间段
     * 3. 将这些时间段的峰值转换为BPM值
     * 4. 将相关性比较高的BPM归化到120BPM左右，然后取众数
     * 
     * @param energySignal 能量信号数组
     * @param sampleRate   采样率
     * @param hopSize      步长
     * @return 计算得到的BPM值
     */
    public double calculateBPM(double[] energySignal, int sampleRate, int hopSize) {
        BPMResult result = calculateBPMWithCount(energySignal, sampleRate, hopSize);
        return result.bpm;
    }

    /**
     * 使用自相关方法计算BPM（改进版），并返回BPM值及其归一化重复次数
     * 1. 计算自相关并找到峰值
     * 2. 计算5-250 BPM对应的时间段
     * 3. 将这些时间段的峰值转换为BPM值
     * 4. 将相关性比较高的BPM归化到120BPM左右，然后取众数
     * 
     * @param energySignal 能量信号数组
     * @param sampleRate   采样率
     * @param hopSize      步长
     * @return BPMResult对象，包含BPM值和归一化重复次数
     */
    public BPMResult calculateBPMWithCount(double[] energySignal, int sampleRate, int hopSize) {
        System.out.println("DEBUG: 开始计算BPM（带计数），能量信号长度 = " + energySignal.length + ", 采样率 = " + sampleRate + ", hopSize = " + hopSize);

        // 添加对空信号或极短信号的检查
        if (energySignal.length < 10) {
            return new BPMResult(0, 0);
        }

        // 检查信号是否全为0
        boolean allZero = true;
        for (double value : energySignal) {
            if (value != 0) {
                allZero = false;
                break;
            }
        }

        if (allZero) {
            return new BPMResult(0, 0);
        }

        // 计算自相关结果
        AutocorrelationResult autocorrResult = computeAutocorrelation(energySignal);
        
        // 找到自相关的峰值（除了0延迟）
        List<Integer> peaks = findPeaks(autocorrResult.autocorrelation);

        if (peaks.isEmpty()) {
            return new BPMResult(0, 0);
        }

        // 对峰值按照自相关值从大到小排序
        peaks.sort((a, b) -> Double.compare(autocorrResult.autocorrelation[b], autocorrResult.autocorrelation[a]));
        
        // 只选择前20%的峰值点
        int topPeakCount = Math.max(1, (int) (peaks.size() * 0.2));
        if (topPeakCount < peaks.size()) {
            peaks = peaks.subList(0, topPeakCount);
        }

        // 计算BPM对应的时间段（使用更精确的计算）
        double maxPeriod = 60.0 / 1; // 12秒
        double minPeriod = 60.0 / 180.0; // 0.24秒
        
        // 将时间段转换为lag值（使用浮点数计算提高精度）
        double maxLag = maxPeriod * sampleRate / hopSize;
        double minLag = minPeriod * sampleRate / hopSize;
        

        // 筛选出在1-180 BPM范围内的峰值，并计算对应的BPM值
        List<BPMCandidate> bpmCandidates = new ArrayList<>();
        for (int peak : peaks) {
            // 使用浮点数计算提高精度
            if (peak >= minLag && peak <= maxLag) {
                double beatPeriod = peak * hopSize / (double) sampleRate;
                double bpm = 60.0 / beatPeriod;
                // 使用样本计数作为额外的权重因子
                double weightedAutocorrelation = autocorrResult.autocorrelation[peak] * Math.sqrt(autocorrResult.sampleCounts[peak] / (double) energySignal.length);
                bpmCandidates.add(new BPMCandidate(peak, bpm, weightedAutocorrelation));
            }
        }
        

        if (bpmCandidates.isEmpty()) {
            return new BPMResult(0, 0);
        }

        // 改进：将相关性比较高的BPM归化到120BPM左右，然后取众数
        return calculateBPMByNormalizationAndModeWithCount(bpmCandidates);
    }

    /**
     * 计算自相关并同时计算归一化因子和样本计数
     * 
     * @param energySignal 能量信号数组
     * @return AutocorrelationResult对象，包含自相关结果、归一化因子和样本计数
     */
    private AutocorrelationResult computeAutocorrelation(double[] energySignal) {
        int n = energySignal.length;
        double[] autocorrelation = new double[n/2];
        double[] normalizationFactors = new double[n/2]; // 添加归一化因子
        int[] sampleCounts = new int[n/2]; // 添加样本计数数组

        // 计算自相关并同时计算归一化因子和样本计数
        for (int lag = 0; lag < n/2; lag++) {
            // 计算信号的局部均值和标准差（用于归一化）
            double sum1 = 0, sum2 = 0, sum1Sq = 0, sum2Sq = 0;
            int count = 0;
            
            for (int i = 0; i < n - lag; i++) {
                double val1 = energySignal[i];
                double val2 = energySignal[i + lag];
                autocorrelation[lag] += val1 * val2;
                sum1 += val1;
                sum2 += val2;
                sum1Sq += val1 * val1;
                sum2Sq += val2 * val2;
                count++;
            }
            
            sampleCounts[lag] = count; // 记录样本数量
            
            // 计算归一化因子（皮尔逊相关系数的分母部分）
            if (count > 0) {
                double mean1 = sum1 / count;
                double mean2 = sum2 / count;
                double var1 = sum1Sq / count - mean1 * mean1;
                double var2 = sum2Sq / count - mean2 * mean2;
                
                // 避免除以零
                if (var1 > 0 && var2 > 0) {
                    normalizationFactors[lag] = Math.sqrt(var1 * var2) * count;
                } else {
                    normalizationFactors[lag] = 1.0; // 默认值
                }
            } else {
                normalizationFactors[lag] = 1.0; // 默认值
            }
            
            // 应用归一化
            if (normalizationFactors[lag] > 0) {
                autocorrelation[lag] /= normalizationFactors[lag];
            }
        }
        
        return new AutocorrelationResult(autocorrelation, normalizationFactors, sampleCounts);
    }

    /**
     * 将相关性比较高的BPM归化到120BPM左右，然后取众数，并返回BPM值和重复次数
     * 
     * @param bpmCandidates BPM候选值列表
     * @return BPMResult对象，包含BPM值和归一化重复次数
     */
    private BPMResult calculateBPMByNormalizationAndModeWithCount(List<BPMCandidate> bpmCandidates) {
        // 1. 过滤掉相关性过低的候选值，只保留前10%相关性较高的
        List<BPMCandidate> filteredCandidates = new ArrayList<>();
        int topCount = Math.max(1, (int) (bpmCandidates.size() * 0.1)); // 至少保留1个
        
        // 按相关性从高到低排序
        bpmCandidates.sort((a, b) -> Double.compare(b.autocorrelationValue, a.autocorrelationValue));
        
        // 只保留前20%的候选值
        for (int i = 0; i < Math.min(topCount, bpmCandidates.size()); i++) {
            filteredCandidates.add(bpmCandidates.get(i));
        }
        

        // 2. 将BPM值归一化到120BPM附近
        double targetBPM = 120.0;
        List<Double> normalizedBPMs = new ArrayList<>();
        
        for (BPMCandidate candidate : filteredCandidates) {
            double originalBPM = candidate.bpm;
            double adjustedBPM = originalBPM;
            double bestDifference = Math.abs(originalBPM - targetBPM);

            // 尝试除以2的幂次来获得更接近targetBPM的值
            double candidateBPM = originalBPM;
            for (int i = 1; i <= 10; i++) { // 限制最多除以2^10
                candidateBPM = candidateBPM / 2.0;
                if (candidateBPM < 20)
                    break; // 避免值过小

                double difference = Math.abs(candidateBPM - targetBPM);
                if (difference < bestDifference) {
                    bestDifference = difference;
                    adjustedBPM = candidateBPM;
                }
            }

            // 也尝试乘以2的幂次来获得更接近targetBPM的值
            candidateBPM = originalBPM;
            for (int i = 1; i <= 10; i++) { // 限制最多乘以2^10
                candidateBPM = candidateBPM * 2.0;
                if (candidateBPM > 400) // 避免值过大
                    break;

                double difference = Math.abs(candidateBPM - targetBPM);
                if (difference < bestDifference) {
                    bestDifference = difference;
                    adjustedBPM = candidateBPM;
                }
            }

            normalizedBPMs.add(adjustedBPM);
        }

        // 3. 对归一化后的BPM值进行分组统计
        Map<Double, Integer> bpmCountMap = new HashMap<>();
        double tolerance = 0.01; // 1%的容差
        
        for (Double normalizedBPM : normalizedBPMs) {
            boolean foundGroup = false;
            // 检查是否与已有的组中的值接近
            for (Double groupBase : bpmCountMap.keySet()) {
                // 计算两个BPM值的比值差异
                double ratio = Math.max(normalizedBPM, groupBase) / Math.min(normalizedBPM, groupBase);
                // 如果比值接近1（差异在1%以内），则认为是同一组
                if (Math.abs(ratio - 1.0) < tolerance) {
                    // 属于同一组，增加计数
                    bpmCountMap.put(groupBase, bpmCountMap.get(groupBase) + 1);
                    foundGroup = true;
                    break;
                }
            }

            // 如果没有找到相近的组，则创建新组
            if (!foundGroup) {
                bpmCountMap.put(normalizedBPM, 1);
            }
        }

        // 4. 改进：当存在多个重复次数相同的BPM组时，统计所有相近BPM的总次数再做对比
        double mostFrequentBPM = 0;
        int maxCount = 0;

        // 遍历所有BPM组，找出具有最高总次数的组
        for (Map.Entry<Double, Integer> entry : bpmCountMap.entrySet()) {
            double currentBPM = entry.getKey();
            int currentCount = entry.getValue();
            
            // 如果当前组的次数更高，或者次数相同但BPM值更小，则选择当前组
            if (currentCount > maxCount || (currentCount == maxCount && currentBPM < mostFrequentBPM)) {
                maxCount = currentCount;
                mostFrequentBPM = currentBPM;
            }
        }


        // 5. 如果没有找到众数，返回相关性最高的BPM
        if (maxCount == 0 || mostFrequentBPM == 0) {
            return new BPMResult(bpmCandidates.get(0).bpm, 1);
        }

        return new BPMResult(mostFrequentBPM, maxCount);
    }



    /**
     * BPM候选类
     */
    private static class BPMCandidate {
        final int lag;
        final double bpm;
        final double autocorrelationValue;
        
        BPMCandidate(int lag, double bpm, double autocorrelationValue) {
            this.lag = lag;
            this.bpm = bpm;
            this.autocorrelationValue = autocorrelationValue;
        }
    }
    
    /**
     * 自相关计算结果类
     */
    private static class AutocorrelationResult {
        final double[] autocorrelation;
        final double[] normalizationFactors;
        final int[] sampleCounts;
        
        AutocorrelationResult(double[] autocorrelation, double[] normalizationFactors, int[] sampleCounts) {
            this.autocorrelation = autocorrelation;
            this.normalizationFactors = normalizationFactors;
            this.sampleCounts = sampleCounts;
        }
    }

    /**
     * BPM结果类，包含BPM值和归一化重复次数
     */
    public static class BPMResult {
        public final double bpm;
        public final int count;
        
        public BPMResult(double bpm, int count) {
            this.bpm = bpm;
            this.count = count;
        }
    }

    /**
     * 查找自相关函数中的峰值
     * 
     * @param data 数据数组
     * @return 峰值位置列表
     */
    private List<Integer> findPeaks(double[] data) {
        List<Integer> peaks = new ArrayList<>();
        // 使用一个较小的阈值来检测峰值，而不是硬性要求大于0
        double threshold = 0.0;
        
        // 计算数据的平均值和标准差，用于动态确定阈值
        double sum = 0;
        for (double value : data) {
            sum += value;
        }
        double mean = data.length > 0 ? sum / data.length : 0;
        
        double sumSq = 0;
        for (double value : data) {
            sumSq += (value - mean) * (value - mean);
        }
        double stdDev = data.length > 1 ? Math.sqrt(sumSq / (data.length - 1)) : 0;
        
        // 设置阈值为均值加上0.1倍标准差，这样可以检测到更多有意义的峰值
        threshold = mean + 0.1 * stdDev;
        
        for (int i = 1; i < data.length - 1; i++) {
            // 检查是否为局部最大值，且大于动态阈值
            if (data[i] > data[i-1] && data[i] > data[i+1] && data[i] > threshold) {
                peaks.add(i);
            }
        }
        
        // 如果没有找到峰值，降低阈值要求，只检查是否为局部最大值
        if (peaks.isEmpty()) {
            for (int i = 1; i < data.length - 1; i++) {
                // 检查是否为局部最大值
                if (data[i] > data[i-1] && data[i] > data[i+1]) {
                    peaks.add(i);
                }
            }
        }
        
        return peaks;
    }
}